package com.crane.common.framework.shiro;

import org.apache.shiro.session.ExpiredSessionException;
import org.apache.shiro.session.InvalidSessionException;
import org.apache.shiro.session.Session;
import org.apache.shiro.session.mgt.SessionContext;
import org.apache.shiro.session.mgt.SessionKey;
import org.apache.shiro.web.servlet.ShiroHttpServletRequest;
import org.apache.shiro.web.session.mgt.DefaultWebSessionManager;
import org.apache.shiro.web.session.mgt.WebSessionManager;
import org.apache.shiro.web.util.WebUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.Serializable;

public class HeaderWebSessionManager extends DefaultWebSessionManager implements WebSessionManager {

    private static final Logger log = LoggerFactory.getLogger(HeaderWebSessionManager.class);

    private boolean sessionIdHeaderEnabled;

    public static final String HEADER_SESSION_NAME = "X-AUTH-SESSION";

    public static final String HEADER_SESSION_ID_SOURCE = "header";

    public HeaderWebSessionManager() {
        super();
        super.setSessionIdUrlRewritingEnabled(false);
        sessionIdHeaderEnabled = true;
    }

    @Override
    protected Session createExposedSession(Session session, SessionContext context) {
        return super.createExposedSession(session, context);
    }

    @Override
    protected Session createExposedSession(Session session, SessionKey key) {
        return super.createExposedSession(session, key);
    }

    @Override
    protected void onStart(Session session, SessionContext context) {
        super.onStart(session, context);

        HttpServletRequest request = WebUtils.getHttpRequest(context);
        HttpServletResponse response = WebUtils.getHttpResponse(context);
        //用header存储会话sessionId
        if (sessionIdHeaderEnabled) {
            Serializable sessionId = session.getId();
            storeSessionId(sessionId, request, response);
        }
    }

    @Override
    public Serializable getSessionId(SessionKey key) {
        Serializable id = super.getSessionId(key);
        if (null != id) {
            return id;
        }
        if (WebUtils.isWeb(key)) {
            ServletRequest request = WebUtils.getRequest(key);
            ServletResponse response = WebUtils.getResponse(key);
            id = getHeaderSessionId(request, response);
        }
        return id;
    }


    protected Serializable getHeaderSessionId(ServletRequest request, ServletResponse response) {
        String id = getHeaderSessionId((HttpServletRequest) request);
        if (id != null) {
            request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_SOURCE, HeaderWebSessionManager.HEADER_SESSION_ID_SOURCE);
        }
        if (id != null) {
            request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID, id);
            request.setAttribute(ShiroHttpServletRequest.REFERENCED_SESSION_ID_IS_VALID, Boolean.TRUE);
        }
        request.setAttribute(ShiroHttpServletRequest.SESSION_ID_URL_REWRITING_ENABLED, isSessionIdUrlRewritingEnabled());
        return id;
    }

    @Override
    protected void onExpiration(Session s, ExpiredSessionException ese, SessionKey key) {
        super.onExpiration(s, ese, key);
        onInvalidation(key);
    }

    @Override
    protected void onInvalidation(Session session, InvalidSessionException ise, SessionKey key) {
        super.onInvalidation(session, ise, key);
        onInvalidation(key);
    }

    @Override
    protected void onStop(Session session, SessionKey key) {
        super.onStop(session, key);
        onInvalidation(key);
    }

    @Override
    public boolean isServletContainerSessions() {
        return false;
    }

    private void storeSessionId(Serializable currentId, HttpServletRequest request, HttpServletResponse response) {
        if (currentId == null) {
            String msg = "sessionId cannot be null when persisting for subsequent requests.";
            throw new IllegalArgumentException(msg);
        }
        response.setHeader(HeaderWebSessionManager.HEADER_SESSION_NAME, currentId.toString());
        log.trace("Set session ID  header for session with id {}", currentId.toString());
    }

    /**
     * 获取请求的token
     */
    private String getHeaderSessionId(HttpServletRequest httpRequest) {
        String sessionId = null;
        //从header中获取session
        sessionId = httpRequest.getHeader(HeaderWebSessionManager.HEADER_SESSION_NAME);
        //如果header中不存在token，则从参数中获取token
        if (!StringUtils.hasText(sessionId)) {
            sessionId = httpRequest.getParameter(HeaderWebSessionManager.HEADER_SESSION_NAME);
        }
        return sessionId;
    }

    private void onInvalidation(SessionKey key) {
        if (WebUtils.isWeb(key)) {
            ServletResponse response = WebUtils.getResponse(key);
            if (null != response) {
                ((HttpServletResponse) response).addHeader(HeaderWebSessionManager.HEADER_SESSION_NAME, "");
            }
        }
    }

    public boolean isSessionIdHeaderEnabled() {
        return sessionIdHeaderEnabled;
    }

    public void setSessionIdHeaderEnabled(boolean sessionIdHeaderEnabled) {
        this.sessionIdHeaderEnabled = sessionIdHeaderEnabled;
    }
}
